"""
数据集操作方法
"""
import os
import random
import cv2
import numpy as np
from skimage import exposure
from image_classification import config


OUTPUT_LIST = config.OUTPUT_LIST
DATASET_DIR = 'datasets/imgs/'


# 将数据集分割为训练数据集，测试数据集和交叉验证集三部分
def split_data(data_dir, output_list, ratio=0.7):
    """
    :param data_dir: 原始图像数据集的路径
    :param output_list: 获取原始图像的标签及其路径，将其分割为训练数据集，测试数据集和交叉验证数据集三部分，此为输出文件的路径
    :param ratio: 数据集分割比例
    :return: 无
    """
    # 读取所有图片的路径
    data_file = os.listdir(data_dir)
    data_list = []  # 存放图片路径
    label_list = []     # 存放图片标签
    for labels in data_file:
        img_file = os.listdir(data_dir + labels)
        for img in img_file:
            img_dir = os.path.join(data_dir, labels, img)
            data_list.append(img_dir)
            label_list.append(labels)
    permutation = np.random.permutation(len(data_list))     # 获取一个乱序的序列
    # 将打乱后的路径写入文件
    with open(output_list[0] + '.txt', 'w') as train:
        for data in permutation[0: int(ratio * len(data_list))]:
            train.write(data_list[data] + '\n')
    with open(output_list[1] + '.txt', 'w') as test:
        for data in permutation[int(ratio * len(data_list)): int((1 + ratio) / 2 * len(data_list))]:
            test.write(data_list[data] + '\n')
    with open(output_list[2] + '.txt', 'w') as val:
        for data in permutation[int((1 + ratio) / 2 * len(data_list)): len(data_list)]:
            val.write(data_list[data] + '\n')
    # 将打乱后的标签写入文件
    with open(output_list[0] + '_label.txt', 'w') as train:
        for label in permutation[0: int(ratio * len(label_list))]:
            train.write(label_list[label] + '\n')
    with open(output_list[1] + '_label.txt', 'w') as test:
        for label in permutation[int(ratio * len(label_list)): int((1 + ratio) / 2 * len(label_list))]:
            test.write(label_list[label] + '\n')
    with open(output_list[2] + '_label.txt', 'w') as val:
        for label in permutation[int((1 + ratio) / 2 * len(label_list)): len(label_list)]:
            val.write(label_list[label] + '\n')


# 从txt中读取数据集
def read_data(data_file, size, augmentation=False):
    """
    :param data_file: 文件路径
    :param size: 输出图像尺寸
    :param augmentation: 是否进行数据增强。数据集样本少及显存小时不应使用数据增强，梯度会难以下降
    :return: 字典类型，第一项为四维图像数组，第二项为二维标签数组，第三项为分类数
    """
    data_list = []
    label_list = []
    classes = []
    train_data = open(data_file + '.txt', 'rU').readlines()
    train_label = open(data_file + '_label.txt', 'rU').readlines()
    print(train_data)
    for i in range(len(train_data)):
        line_data = train_data[i].strip('\n')   # 取出行末的换行符
        line_label = train_label[i].strip('\n')
        if line_label not in classes:
            classes.append(line_label)
        img = cv2.imread(line_data)
        img = cv2.resize(img, size)
        # 进行数据增强，得到两张图片
        if augmentation:
            aug1 = data_augmentation(img)
            aug2 = data_augmentation(img)
            data_list.append(aug1)
            label_list.append(line_label)
            data_list.append(aug2)
            label_list.append(line_label)
        else:
            data_list.append(img)
            label_list.append(line_label)
    return data_list, label_list, len(classes)


# 对图像进行数据增强
def data_augmentation(img, flip=True, transpose=True, gamma=True, log=True):
    """
    :param img: 输入图像
    :param flip: 是否进行随机反转
    :param transpose: 是否进行随机旋转
    :param gamma: 是否进行随机亮度值调节
    :param log: 是否进行随机对比度调节
    :return: 输出图像
    """
    # 随机反转
    if flip:
        random_flip = random.randint(-1, 1)
        img = cv2.flip(img, random_flip)
    # 随机旋转
    if transpose:
        random_transpose = random.randint(0, 1)
        if random_transpose == 0:
            img = cv2.transpose(img)
    # 随机亮度
    if gamma:
        random_gamma = round(random.uniform(0.5, 1.5), 1)
        img = exposure.adjust_gamma(img, random_gamma)
    # 随机对比度
    if log:
        random_log = round(random.uniform(0.75, 1.0), 1)
        img = exposure.adjust_log(img, random_log)
    return img


if __name__ == '__main__':
    split_data(DATASET_DIR, OUTPUT_LIST)
    print('Processing Completed!')
